package com.yunmaozj.tools.lock.web.interceptor;

import com.yunmaozj.tools.lock.RedisLock;
import com.yunmaozj.tools.lock.web.annotations.RequestValve;
import lombok.extern.slf4j.Slf4j;
import org.springframework.util.StringUtils;
import org.springframework.web.method.HandlerMethod;
import org.springframework.web.servlet.HandlerInterceptor;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.io.IOException;
import java.util.Map;

/**
 * 请求限制拦截器
 *
 * @author zhongminghong
 * @email zhongmh@yunmaozj.com
 * @create 2020-09-08 9:25
 */
@Slf4j
public class RequestInterceptor implements HandlerInterceptor {

    private RedisLock redisLock = new RedisLock("localhost", 6379, RedisLock.LockType.VALVE);

    @Override
    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) throws Exception {
        return tryLock(request, response, (HandlerMethod) handler);
    }

    @Override
    public void afterCompletion(HttpServletRequest request, HttpServletResponse response, Object handler, Exception ex) throws Exception {

    }

    /**
     * @param request
     * @param response
     * @param handler
     * @return
     */
    private boolean tryLock(HttpServletRequest request, HttpServletResponse response, HandlerMethod handler) {
        RequestValve valve = handler.getMethodAnnotation(RequestValve.class);

        if (valve != null) {
            StringBuilder cacheKeyBuilder = new StringBuilder("limit_")
                    .append(handler.getMethod().getDeclaringClass().getSimpleName()).append(".")
                    .append(handler.getMethod().getName()).append("(");
            StringBuilder parameterBuilder = new StringBuilder();
            int parameterNumber = 0;
            for (String paramName : valve.parameters()) {
                String value = request.getParameter(paramName);
                if (StringUtils.hasText(value)) {
                    parameterBuilder.append(value).append(",");
                    parameterNumber++;
                }
            }
            for (String headerName : valve.headers()) {
                String value = request.getHeader(headerName);
                if (StringUtils.hasText(value)) {
                    parameterBuilder.append(value).append(",");
                    parameterNumber++;
                }
            }
            if (parameterNumber != 0 && parameterBuilder.length() > 0) {
                parameterBuilder.deleteCharAt(parameterBuilder.length() - 1);
            } else {
                // 没有需要限流的参数
                log.warn("not find request limit parameter of {}", cacheKeyBuilder + ")");
                return true;
            }
            //打印原生参数
            if (log.isDebugEnabled())
                log.debug("request limit string {}", cacheKeyBuilder + parameterBuilder.toString() + ")");
            cacheKeyBuilder.append(parameterBuilder.toString().hashCode()).append(")");
            String cacheKey = cacheKeyBuilder.toString();
            if (log.isDebugEnabled())
                log.debug("request limit cache key {}", cacheKey);
            if (valve.policy() == RequestValve.LimitPolicy.REJECT) {
                if (!redisLock.tryLock(cacheKey)) {
                    try {
                        response.setCharacterEncoding("UTF-8");
                        response.getWriter().write(valve.responseMessage());
                        return false;
                    } catch (IOException e) {
                        e.printStackTrace();
                    }
                }
            } else {
                redisLock.lock(cacheKey);
            }

        }
        return true;
    }
}
